#include "Improved.hpp"

Improved::Improved(yjy::Config *cf)
{
    this->cf = cf;
    conf_last = 1;
    cadidator_num = this->cf->getICF().candidatorNum;
}

Improved::~Improved()
{
    //todo
    cv::COLOR_BGR2Lab;
    cv::COLOR_Lab2BGR;
}

int Improved::calculateWS(float confidence, float data, int windowSize)
{
    //todo
    if (!this->cf->getICF().dynamicWS)
        return windowSize;

    int dws;
    float weight_alpha = 0.99f, weight_beta = 1 - weight_alpha;
    windowSize *= 1.5f;

    float f_xy = 0.63662f * acos(data);

    dws = f_xy * windowSize * weight_alpha + windowSize * weight_beta + 0.5f;
    return dws % 2 ? dws : dws + 1;
}

float Improved::calculatePriority(float data_term, float conf_term)
{
    if (!this->cf->getICF().new_priority)
        return data_term * conf_term;
    //3.7979f
    return (4.0f) * data_term * conf_term + (0) * conf_term + (1) * data_term;
}

void Improved::cvtColor(cv::Mat &origFloat, cv::Mat &origCIE)
{ //44
    if (this->cf->getICF().CIE_color != -1)
        cv::cvtColor(origFloat, origCIE, this->cf->getICF().CIE_color);
    else
        origCIE = origFloat;
}

void Improved::cvtColor_inverse(cv::Mat &orig, cv::Mat &dst)
{ //56
    if (this->cf->getICF().CIE_color_inverse != -1)
        cv::cvtColor(orig, dst, this->cf->getICF().CIE_color_inverse);
    else
        dst = orig;
}

void Improved::visualize(cv::Mat &image_padded, cv::Mat &mask_padded, int windowSize)
{
    cv::imshow("result_improved", image_padded);
}

void Improved::precedure(cv::Mat &image_padded, cv::Mat &mask_padded)
{
    cv::inpaint(image_padded, mask_padded, image_padded, inpainting_radius, cv::INPAINT_TELEA);
}

void Improved::findCandidatePatch(cv::Mat &image_padded, cv::Mat &mask_padded, int i, int j, int windowSize, std::vector<std::pair<float, cv::Point2i>> &candidators)
{
    const float MAXFLOAT = 3.40282347e+38F;
    int offset = windowSize / 2;
    float min_ssd = MAXFLOAT, ssd;

    // get the patch around the target
    cv::Range targetRowRange_padded(i - offset, i + offset + 1);
    cv::Range targetColRange_padded(j - offset, j + offset + 1);

    // make a CIE color version of input
    cv::Mat image_cie_padded;
    cvtColor(image_padded, image_cie_padded);

    cv::Mat imageTemplate = image_cie_padded.rowRange(targetRowRange_padded).colRange(targetColRange_padded);
    cv::Mat maskTemplate = mask_padded.rowRange(targetRowRange_padded).colRange(targetColRange_padded);
    cv::Mat maskTemplateInv, image_padded_gray, dx, dy;

    cv::cvtColor(image_padded, image_padded_gray, cv::COLOR_BGR2GRAY);
    /*
    cv::Sobel(image_padded_gray.rowRange(targetRowRange_padded).colRange(targetColRange_padded), dx, CV_32F, 0, 1, 3, 1, 0, cv::BORDER_CONSTANT);
    cv::Sobel(image_padded_gray.rowRange(targetRowRange_padded).colRange(targetColRange_padded), dy, CV_32F, 1, 0, 3, -1, 0, cv::BORDER_CONSTANT);
    cv::Point2i dxdy = cv::Point2i(dx.at<float>(dx.cols / 2, dx.rows / 2), dy.at<float>(dy.cols / 2, dy.rows / 2));
    */
    cv::Sobel(image_padded_gray, dx, CV_32F, 0, 1, 3, 1, 0, cv::BORDER_CONSTANT);
    cv::Sobel(image_padded_gray, dy, CV_32F, 1, 0, 3, -1, 0, cv::BORDER_CONSTANT);
    cv::Point2f dxdy = max_edge_vector(dx, dy, i, j, offset);

    float dz = abs_vector(dxdy.x, dxdy.y);
    if (this->cf->getICF().isQuikMatch && dz < dz_threshold)
    {
        printf("find a very small dz,take it as the best matching one!\n");
        candidators.push_back(std::pair<float, cv::Point2i>(0, cv::Point2i(i, j)));
        return;
    }

    maskTemplate.convertTo(maskTemplateInv, CV_8U, -1, 1);
    cv::Mat candidSourcePatch;
    std::map<float, cv::Point2d> m;

    for (int k = windowSize; k < image_padded.rows - windowSize; ++k)
    {
        for (int l = windowSize; l < image_padded.cols - windowSize; ++l)
        {
            if (k - offset < windowSize / 2 + 1 || k + offset + 1 > image_padded.rows - windowSize / 2 - 1 || l - offset < windowSize / 2 + 1 || l + offset + 1 > image_padded.cols - windowSize / 2 - 1)
            {
                continue;
            }
            //check that this candidate source patch does not have overlap with target patch
            if (abs(k - i) < imageTemplate.rows || abs(l - j) < imageTemplate.cols)
            {
                // target and source patch overlap
                continue;
            }
            cv::Range candidSourceRowRange = cv::Range(k - offset, k + offset + 1);
            cv::Range candidSourceColRange = cv::Range(l - offset, l + offset + 1);

            cv::Scalar s = cv::sum(mask_padded.rowRange(candidSourceRowRange)
                                       .colRange(candidSourceColRange));
            if (s[0] > 0)
            {
                // target region overlaps with masked region
                continue;
            }
            candidSourcePatch = image_cie_padded.rowRange(candidSourceRowRange).colRange(candidSourceColRange);

            ssd = cv::norm(imageTemplate, candidSourcePatch, cv::NORM_L2, maskTemplateInv);

            m.insert(std::pair<float, cv::Point2d>(ssd, cv::Point2i(k, l)));
        }
    }

    std::map<float, cv::Point2d>::iterator iter = m.begin();
    for (int count = 0; count < cadidator_num; count++, ++iter)
    {
        candidators.push_back(*iter);
    }
}

cv::Point2i Improved::findBestPatch(cv::Mat &image_padded, cv::Mat &mask_padded, int i, int j, int windowSize, std::vector<std::pair<float, cv::Point2i>> &candidators, cv::Point2f dxdy)
{
    const float MAXFLOAT = 3.40282347e+38F;

    float min_diff = MAXFLOAT, max_diff = 0;

    int i_min = -1, j_min = -1, i_max = -1, j_max = -1;
    int offset = windowSize / 2;

    cv::Mat image_padded_gray, maskTemplate;

    cv::Range targetRowRange_padded(i - offset, i + offset + 1);
    cv::Range targetColRange_padded(j - offset, j + offset + 1);

    cv::Range sourceRowRange_padded;
    cv::Range sourceColRange_padded;

    maskTemplate = mask_padded.rowRange(targetRowRange_padded).colRange(targetColRange_padded);

    for (std::vector<std::pair<float, cv::Point2i>>::iterator iter = candidators.begin(); iter != candidators.end(); ++iter)
    {
        cv::Mat dx, dy, nx, ny;

        float ssd = iter->first;
        float i_m = iter->second.x;
        float j_m = iter->second.y;

        sourceRowRange_padded = cv::Range(i_m - offset, i_m + offset + 1);
        sourceColRange_padded = cv::Range(j_m - offset, j_m + offset + 1);

        //try patch,than check the new (dy,-dx)'s direction
        image_padded.rowRange(sourceRowRange_padded).colRange(sourceColRange_padded).copyTo(image_padded.rowRange(targetRowRange_padded).colRange(targetColRange_padded), maskTemplate);
        cv::cvtColor(image_padded, image_padded_gray, cv::COLOR_BGR2GRAY);
        cv::Sobel(image_padded_gray, dx, CV_32F, 0, 1, 3, 1, 0, cv::BORDER_CONSTANT);
        cv::Sobel(image_padded_gray, dy, CV_32F, 1, 0, 3, -1, 0, cv::BORDER_CONSTANT);
        dx.setTo(cv::Scalar(0), mask_padded);
        dy.setTo(cv::Scalar(0), mask_padded);

        cv::Point2f new_dxdy = max_edge_vector(dx, dy, i, j, 2 * offset);

        double cos_ab = cos_vector(dxdy, new_dxdy);
        float diff = 1 * ssd * acos(cos_ab) + 5 * ssd;
        if (min_diff > diff)
        {
            i_min = i_m;
            j_min = j_m;
            min_diff = diff;
        }
        if (max_diff < diff)
        {
            i_max = i_m;
            j_max = j_m;
            max_diff = diff;
        }
    }

    //all cadidators are too bad
    if (false) //min_diff > 120
    {
        printf("all cadidators are too bad！\n");
        return cv::Point2i(i, j);
    }
    else
    {
        printf("i_m=%d,j_m=%d,offset=%d\n", i_min, j_min, offset);
        return cv::Point2i(i_min, j_min);
    }
}

void Improved::afterIteration(cv::Mat &image_padded, cv::Mat &mask_padded, cv::Mat &confidence_padded, cv::Mat &data_paddeed, int i, int j)
{
    if (confidence_padded.at<float>(i, j) < conf_last / 2 || confidence_padded.at<float>(i, j) > conf_last * 4)
    {
        conf_last = confidence_padded.at<float>(i, j);
        //inpainting_radius = inpainting_radius / 2;
        cv::inpaint(image_padded, mask_padded, image_padded, inpainting_radius, cv::INPAINT_TELEA);
    }
}
